#include "yolo.h"
using namespace std;


namespace {


const std::vector<std::vector<int>> m_anchors = {{10, 13, 16, 30, 33, 23}, {30, 61, 62 ,45, 59, 119}, {116, 90, 156, 198, 373, 326}};
const int ojb_class_num = 3;
const int cell_size = 8;
const std::vector<int> stride = {8, 16, 32};

const double g_NMSThreshold = 0.45;
const double g_scoreThreshold = 0.25;

 const vector<cv::Scalar> g_colors {
        cv::Scalar(0, 0, 255), cv::Scalar(255, 0, 0), cv::Scalar(0, 125, 255)};
}


yolo::yolo(const char* modelPath, uint32_t modelWidth,
                           uint32_t modelHeight)
    :g_modelWidth_(modelWidth), g_modelHeight_(modelHeight)
{
    g_modelPath_ = modelPath;
    
}

yolo::~yolo()
{
}


Result yolo::Init() {

    // pInstant = new basic (g_modelPath_, g_modelWidth_,  g_modelHeight_);
    Yolo_Instant = new basic(g_modelPath_, g_modelWidth_,  g_modelHeight_);

    Result ret = Yolo_Instant->Init_atlas();
    ret = Yolo_Instant->CreateInput();
    ret = Yolo_Instant->CreateOutput();

    return SUCCESS;
}

Result yolo::UnInit() {
    Yolo_Instant->DestroyResource();
}


Result yolo::Preprocess(unsigned char *frame) {
    
    struct  timeval tstart,tend;
    double timeuse;
    gettimeofday(&tstart,NULL);

    aclError ret = aclrtMemcpy(Yolo_Instant->g_imageDataBuf_, Yolo_Instant->g_imageDataSize_,
                               frame, Yolo_Instant->g_imageDataSize_, ACL_MEMCPY_HOST_TO_DEVICE);
    if (ret) {
        std::cout<<"Copy resized image data to device failed."<<std::endl;
        return FAILED;
    }

    gettimeofday(&tend,NULL);
    timeuse = 1000000*(tend.tv_sec - tstart.tv_sec) + \
				(tend.tv_usec - tstart.tv_usec);
    std::cout<<"Preprocess time:  "<< timeuse/1000<<"ms"<<std::endl;

    return SUCCESS;
}




Result yolo::Inference() {
    
    struct  timeval tstart,tend;
    double timeuse;
    gettimeofday(&tstart,NULL);

    Result ret = Yolo_Instant->inference();

    gettimeofday(&tend,NULL);
    timeuse = 1000000*(tend.tv_sec - tstart.tv_sec) + \
				(tend.tv_usec - tstart.tv_usec);
    std::cout<<"Inference time:  "<< timeuse/1000<<"ms"<<std::endl;

    return SUCCESS;
}


vector<BBox> yolo::Postprocess() {
    struct  timeval tstart,tend,tbuffer;
    double timeuse;
    gettimeofday(&tstart,NULL);

    vector <BBox> boxes;

    for(int n = 0; n < Yolo_Instant->outputSize; n++)//Yolo_Instant->outputSize
    {   
        // n = 0;
        // std::cout<<"步长为："<<stride[n]<<std::endl;

        int grid_h = g_modelHeight_ / stride[n];
        int grid_w = g_modelWidth_ / stride[n];
        // std::cout<<"grid_w = "<<grid_w<<", "<<"grid_h = "<<grid_h<<std::endl;
        int grid_len = grid_w * grid_h;
        float* input = (float *)Yolo_Instant->GetInferenceOutputItem(n);    // 1x(n+5)x3 x80x80
        
        std::vector<int> anchor = m_anchors.at(n);
        for(int a = 0; a < 3; a++){
            // std::cout<<"anchor = "<<anchor[a * 2]<<" "<<anchor[a * 2 +1]<<std::endl;
            for(int i = 0; i < grid_h; i++){
                for(int j = 0; j < grid_w; j++){
                    float box_conf = input[(cell_size * a + 4) * grid_len + i * grid_w + j];
                    // float box_conf = sigmoid(input[(a *  + 4) * grid_len + i * grid_w + j]);
                    if(box_conf > g_scoreThreshold){       // 初筛
                        int     offset = (cell_size* a) * grid_len + i * grid_w + j;
                        float* in_ptr = input + offset;
                        float   box_x  = 2 * (*in_ptr) - 0.5;
                        float   box_y  = 2 * (in_ptr[1 * grid_len]) - 0.5;
                        float   box_w  = (in_ptr[2 * grid_len]);
                        float   box_h  = (in_ptr[3 * grid_len]);
                        box_x          = (box_x + j) * (float)stride[n];
                        box_y          = (box_y + i) * (float)stride[n];
                        box_w          = 4 * box_w * box_w * (float)anchor[a * 2];     //n*6为vector索引
                        box_h          = 4 * box_h * box_h * (float)anchor[a * 2 + 1];
                        
                        float maxClassProbs = in_ptr[5 * grid_len];
                        int   maxClassId    = 0;
                        for (int k = 0; k < ojb_class_num; ++k) {
                            float prob = in_ptr[(5 + k) * grid_len];
                            if (prob > maxClassProbs) {
                                maxClassId    = k;
                                maxClassProbs = prob;
                            }
                        }
        
                        float prob_t = box_conf * maxClassProbs;
                        if(prob_t > g_scoreThreshold){       // 细筛
                            BBox b;
                            b.x = box_x;
                            b.y = box_y;
                            b.w = box_w;
                            b.h = box_h;
                            b.score = box_conf;
                            b.classIndex = maxClassId;
                            boxes.push_back(b);
                        }
                    }
                }
            }
        }

        if (Yolo_Instant->g_runMode_ == ACL_HOST) {
	        aclError ret = aclrtFreeHost(input);
        } 
    }   
       
    vector<BBox> result;
    NMS(boxes, result);

    gettimeofday(&tend,NULL);
    timeuse = 1000000*(tend.tv_sec - tstart.tv_sec) + (tend.tv_usec - tstart.tv_usec);
    std::cout<<"Postprocess time: "<<timeuse/1000<<std::endl;

    return result;
}






void yolo::NMS(vector<BBox> &boxes, vector<BBox> &result)
{
    result.clear();
    sort(boxes.begin(), boxes.end(), SortScore);  //按照置信度从大到小排序

    while (boxes.size() != 0) {
        result.push_back(boxes[0]);
        size_t index = 1;
        while (boxes.size() > index) {
            float iou = IOU(boxes[0], boxes[index]);
            if (iou > g_NMSThreshold) {
                boxes.erase(boxes.begin() + index);
                continue;
            }
            ++index;
            //cout<<index<<endl;
        }
        boxes.erase(boxes.begin());
    }
}

float yolo::IOU(const BBox &b1, const BBox &b2)
{
    float x1 = max(b1.x - b1.w / 2, b2.x - b2.w / 2);
    float y1 = max(b1.y - b1.h / 2, b2.y - b2.h / 2);
    float x2 = min(b1.x + b1.w / 2, b2.x + b2.w / 2);
    float y2 = min(b1.y + b1.h / 2, b2.y + b2.h / 2);
    float w = max(0.0f, x2 - x1 + 1);
    float h = max(0.0f, y2 - y1 + 1);
    float area = w * h;
    return area / (b1.w * b1.h + b2.w * b2.h - area);
}

bool yolo::SortScore(BBox box1, BBox box2)
{
    return box1.score > box2.score;
}

void yolo::DrawBoundBoxToImage(vector<BBox>& detectionResults, cv::Mat& image)
{
    int half = 2;
    float w = float(image.cols) / float(g_modelWidth_);
    float h = float(image.rows) / float(g_modelHeight_);
    for (int i = 0; i < detectionResults.size(); ++i) {
        cv::Point p1, p2;
        p1.x = (uint32_t)((detectionResults[i].x - detectionResults[i].w / half) * w);
        p1.y = (uint32_t)((detectionResults[i].y - detectionResults[i].h / half) * h);
        p2.x = (uint32_t)((detectionResults[i].x + detectionResults[i].w / half) * w);
        p2.y = (uint32_t)((detectionResults[i].y + detectionResults[i].h / half) * h);
        cv::rectangle(image, p1, p2, g_colors[detectionResults[i].classIndex], 2);
        char name[50];
        sprintf(name, "%.1d_%.3f", detectionResults[i].classIndex, detectionResults[i].score);
        // 图片，添加的文字，左上角坐标，字体，字体大小，颜色，字体粗细
        cv::putText(image, name, p1, cv::FONT_HERSHEY_SIMPLEX, 1, g_colors[detectionResults[i].classIndex], 1);
	}

    // cv::imshow("output_pic", image);
    cv::imwrite("output_pic.jpg", image);
}
